import os
import json
import sys
import numpy as np
import cv2
import re
import imageio


class DataEnhance:

    def __init__(self, data_dir, data_dict, flag=False):
        self.flag = flag
        self.object_count = 0
        self.label_info = []
        self.region_info = []
        self.image_path = data_dir + os.sep + data_dict["filename"]
    
        if not os.path.exists(self.image_path):
            raise Exception("file not exist", 300)
        if not self.flag:
            self.ori_image = cv2.imread(self.image_path, 0)
            
        else:
            self.ori_image = imageio.imread(self.image_path)
        self.save_image = self.ori_image.copy()
        self._load_data(data_dict)
        self._fill_poly_all()

    def _load_data(self, data_dict):
        for one_label in data_dict["regions"]:
            if not self.flag:
                label = one_label["region_attributes"]["type"]
            else:
                label = one_label["region_attributes"]["type"] + "_" + one_label["region_attributes"]["isliving"]
            xs = one_label["shape_attributes"]["all_points_x"]
            ys = one_label["shape_attributes"]["all_points_y"]
            xy = []
            for x,y in zip(xs,ys):
                xy.append([x,y])
            
            self.label_info.append(label)
            self.region_info.append(xy)

    def convert_to_dict(self, image_path):
        imageio.imwrite(image_path, self.save_image)

        return_data = dict()
        return_data["filename"] = os.path.basename(image_path)
        regions = []
        for idx in range(len(self.label_info)):
            region = dict()
            if not self.flag:
                region["region_attributes"] = {"type": self.label_info[idx]}
            else:
                region["region_attributes"] = {"type": self.label_info[idx].split("_")[0], 
                                                "isliving":self.label_info[idx].split("_")[1]}
            region["shape_attributes"] = {"name": "polygon",
                                            "all_points_x": [int(x[0]) for x in self.region_info[idx]],
                                            "all_points_y": [int(x[1]) for x in self.region_info[idx]]}
            regions.append(region)
        return_data["regions"] = regions
        return_data["file_attributes"] = {}
        return return_data

    def random_add_region(self, image, label, region):
        mask_big = np.zeros(image.shape[:2], dtype=np.uint8)
        # print("mask_big: ", mask_big.shape)
        # print("region: ", region)
        mask_big = cv2.fillConvexPoly(mask_big, region, [255])

        bound_rect = cv2.boundingRect(region)
        mask_small = mask_big[bound_rect[1]:bound_rect[1] + bound_rect[3], 
                                bound_rect[0]:bound_rect[0] + bound_rect[2]].copy()

        # get mask in func
        new_rect,mask_small_new  = self._get_random_rect(mask_small, 
                                                            True if self.flag and (label== "F_living" or label=="F_dead") else False)
        if new_rect is None:
            print("can't find a rect")
            return

        mask_background = self.save_image[new_rect[1]:new_rect[1] + new_rect[3], 
                                new_rect[0]:new_rect[0] + new_rect[2]]
        mask_front = image[bound_rect[1]:bound_rect[1] + bound_rect[3], 
                                bound_rect[0]:bound_rect[0] + bound_rect[2]].copy()

        if not self.flag:
            # only add text mask
            mask_small_tem = cv2.threshold(mask_front, 120, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
            mask_small_new = mask_small_tem[1]
        #print(mask_small_new)
        cv2.copyTo(mask_front, mask_small_new, mask_background)
        
        # update
        offset_x = -bound_rect[0] + new_rect[0]
        offset_y = -bound_rect[1] + new_rect[1]
        new_region = []
        for x,y in region:
            new_region.append([x + offset_x, y + offset_y])
        new_region = np.array(new_region)
        self._update_data(label, new_region)


    def _get_random_rect(self, mask_small, small_flag=False):
            
        height, width = self.ori_image.shape[:2]
        mask_height, mask_width = mask_small.shape[:2]
        if not self.flag:
            choice_y = [x for x in range(0,height-mask_height, 20)]
            choice_x = [x for x in range(0,width-mask_width, 20)]
        else:
            choice_y = [x for x in range(30,height-mask_height - 30, 10)]
            choice_x = [x for x in range(200,width-mask_width - 200, 10)]

        count = 0
        
        new_rect = 0
        posi = 0

        while count < 100:
            count += 1
            tmp_y = np.random.choice(choice_y)
            tmp_x = np.random.choice(choice_x)

            ori_mask = self.all_poly_mask[tmp_y:tmp_y+mask_height, tmp_x:tmp_x+mask_width].copy() 
            std_mask = np.zeros(self.ori_image.shape[:2], dtype=np.uint8)
            std_mask[tmp_y:tmp_y+mask_height, tmp_x:tmp_x+mask_width] = mask_small
        
            if np.sum(ori_mask) != 0 and not small_flag:
                first_whole_mask = self.get_whole_mask( tmp_x, tmp_y, mask_height, mask_width, std_mask)
                std_mask_sum = np.sum(std_mask)
                both_have = cv2.bitwise_and(std_mask, first_whole_mask)
                both_have_sum = np.sum(both_have)
                if both_have_sum / std_mask_sum < 0.35: # 有
                    new_rect = [tmp_x, tmp_y, mask_width, mask_height]
                    posi = (std_mask - both_have)[tmp_y:tmp_y+mask_height, tmp_x:tmp_x+mask_width]
                    break
            else:
                new_rect = [tmp_x, tmp_y, mask_width, mask_height]
                posi = mask_small
                break

        if type(new_rect) == int:
            return None, None
        return new_rect, posi

    def get_whole_mask(self, tmp_x, tmp_y, mask_height, mask_width, std_mask):
        whole_mask = np.zeros(self.ori_image.shape[:2], dtype=np.uint8)

        for region in self.region_info:
            new_mask = np.zeros(self.ori_image.shape[:2], dtype=np.uint8)
            # print("new_mask.shape: ", new_mask.shape)
            # print("region: ", region)
            new_mask = cv2.fillConvexPoly(new_mask, np.array(region), [255])
            out_mask = cv2.bitwise_and(std_mask, new_mask)
            if np.sum(out_mask) > 2:
                whole_mask = cv2.fillConvexPoly(whole_mask, np.array(region), [255])
        
        return whole_mask

    def _get_random_rect_cover(self):
        pass

    def _fill_poly_all(self):
        self.all_poly_mask = np.zeros(self.ori_image.shape[:2], dtype=np.uint8)
        self.object_count = len(self.region_info)
        for region in self.region_info:
            region = np.array(region)
            self.all_poly_mask = cv2.fillConvexPoly(self.all_poly_mask, region, [255])

    def _update_data(self, label, region):
        self.label_info.append(label)
        self.region_info.append(region)
        self.all_poly_mask = cv2.fillConvexPoly(self.all_poly_mask, region, [255])
